import os
import tempfile
import uuid
from PyPDF4 import PdfFileReader

from pdf2image import convert_from_bytes
from PIL import Image

POPPLER_PATH = r'C:\Users\YANG.LEI\Downloads\poppler-0.68.0_x86\poppler-0.68.0\bin'

process_pages_each_time = 10  # 每次转换多少页，太多的话内存不够
out_d = 'jpg'
clean_d = 'clean_jpg'
os.makedirs(out_d, exist_ok=True)
os.makedirs(clean_d, exist_ok=True)


def get_pages_count(file):
    pdf = PdfFileReader(open(file, 'rb'))
    return pdf.getNumPages()


def split_page_parts(total_count, each_count):
    all_parts = []
    for i in range(0, len(total_count), each_count):
        all_parts.append(total_count[i:i + each_count])
    return all_parts


def generate_page_jpg_name(page_num):
    return f'{out_d}/{page_num}.jpg'


def extract_pages(pdf_bytes, page_nums):
    first_page = page_nums[0]
    last_page = page_nums[-1]
    if all(os.path.isfile(generate_page_jpg_name(p)) for p in page_nums):
        return
    print(f'converting page {first_page} - {last_page}')
    with tempfile.TemporaryDirectory() as tmpdirname:
        pages = convert_from_bytes(pdf_bytes, 300,
                                   first_page=first_page,
                                   last_page=last_page,
                                   poppler_path=POPPLER_PATH,
                                   output_folder=tmpdirname)
        for page_num, page in zip(page_nums, pages):
            page.save(generate_page_jpg_name(page_num), 'JPEG')


def crop_images():
    try:
        for img_file in os.listdir(out_d):
            src_full = os.path.join(out_d, img_file)
            out_file = os.path.join(clean_d, img_file)
            if os.path.isfile(out_file):
                continue
            print(src_full)
            img = Image.open(src_full)
            # img.show()
            w, h = img.size
            box = (0, 80, w, h)
            croppedImage = img.crop(box)
            croppedImage.save(out_file)
            # croppedImage.show()

    except Exception as e:
        print(e)


def convert_pdf_to_images(pdf_file):
    pdf_bytes = open(pdf_file, 'rb').read()
    pages_count = get_pages_count(pdf_file)
    print(f'{pdf_file} total pages = {pages_count}')
    for part in split_page_parts(range(1, pages_count + 1), process_pages_each_time):
        extract_pages(pdf_bytes, part)


if __name__ == "__main__":
    pdf_file = r"C:\Users\YANG.LEI\Downloads\蒋中正先生手改荒漠甘泉译稿.pdf"
    convert_pdf_to_images(pdf_file)
    crop_images()
